/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 * Description: TbeVectorReduce类的实现文件
 * Author: jiasonghua
 * Create: 2023-2-3
 */

#include "tbe_vector_reduce.h"
#include "elewise_memset.h"
#include "mem_device_pub.h"
#include <fstream>
#include <dlfcn.h>
#include <unistd.h>
#include <nlohmann/json.hpp>
#include <cstring>
#include <hccl/hccl_types.h>
#include "runtime/rt.h"
#include "op_json_info.h"
#include "op_tiling.h"
#include "vector_tiling.h"
#include "eletwise_v2.h"
#include "adapter_rts.h"
#include "mmpa_api.h"
#include "tbe_crack_cleared.h"

namespace TbeReduce {
using namespace std;

TbeCrackCleard::TbeCrackCleard()
    : addrListDevPtr_(nullptr), initTilingDataHostPtr_(false)
{}

TbeCrackCleard::~TbeCrackCleard()
{
    if (addrListDevPtr_ != nullptr) {
        if (hrtFree(addrListDevPtr_) != HCCL_SUCCESS) {
            HCCL_WARNING("free device memory failed");
        }
        addrListDevPtr_ = nullptr;
    }
}

HcclResult TbeCrackCleard::Init()
{
    CHK_PRT_RET(isInit_, HCCL_DEBUG("TbeVectorReduce had been initialized"), HCCL_SUCCESS);
    CHK_RET(hrtGetDeviceType(deviceType_));
    switch (deviceType_) {
        case DevType::DEV_TYPE_910:
            InitOpInfoMap910A();
            break;
        case DevType::DEV_TYPE_310P3:
            InitOpInfoMap310P3();
            break;
        case DevType::DEV_TYPE_910B:
            InitOpInfoMap910B();
            break;
        case DevType::DEV_TYPE_910_93:
            break;
        default:
            HCCL_ERROR("[TbeCrackCleard][Init]devType[%u] is not supported", deviceType_);
            return HCCL_E_NOT_SUPPORT;
    }

    CHK_RET(LoadOpBinary());

    isInit_ = true;

    return HCCL_SUCCESS;
}

HcclResult TbeCrackCleard::GetTilingRunInfo(nlohmann::json &opDescInfo, nlohmann::json &opTilingInfo,\
    OpRunInfo &runInfo, std::vector<std::int64_t> &crackSize)
{
    MemSetCompileInfo CompileInfo;
    CompileInfo.workspace_num = opTilingInfo["vars"]["workspace_num"];
    CompileInfo.ub_size = opTilingInfo["vars"]["ub_size"];
    CompileInfo.core_num = opTilingInfo["vars"]["core_num"];
    CompileInfo.max_repeat_time = opTilingInfo["vars"]["max_repeat_time"];

    std::string opName;
    CHK_RET(GetOpProperty(opDescInfo, "binFileName", opName));

    bool isTiling = MemSetTiling(opName, crackSize, CompileInfo, runInfo);
    CHK_PRT_RET(!isTiling, HCCL_ERROR("[Get][TilingRunInfo]crack cleard tiling failed"), HCCL_E_CCE);

    return HCCL_SUCCESS;
}

HcclResult TbeCrackCleard::Run(std::vector<std::int64_t> &crackAddr, std::vector<std::int64_t> &crackSize,
    HcclRtStream stream)
{
    if (deviceType_ == DevType::DEV_TYPE_910_93) { // 当前tensor间缝隙清零支持910 910B 310p3
            HCCL_WARNING("[TbeCrackCleard][Run] not support 910_93");
            return HCCL_SUCCESS;
        }

    if (!initTilingDataHostPtr_) {
        tilingDataHostPtr_ = new (std::nothrow) char[(crackAddr.size() + 1) * MEMSET_CRACK_TILING_DATA_MAX_SZIE];
        CHK_PTR_NULL(tilingDataHostPtr_);
    }

    // 获取op关联的描述信息和tiling信息
    nlohmann::json opDescInfo;
    nlohmann::json opTilingInfo;
    CHK_RET(GetOpInfo(opDescInfo, opTilingInfo));

    // 获取tiling data
    OpRunInfo runInfo;
    runInfo.blockDim = 0;
    runInfo.workspaces = { 0 };
    runInfo.clearAtomic = 0;
    CHK_RET(GetTilingRunInfo(opDescInfo, opTilingInfo, runInfo, crackSize));

    // 将tilingdata拷贝device侧
    std::unique_lock<std::mutex> tilingDataLock(tilingDataMutex_);
    void *tilingDataDevPtr = nullptr;
    CHK_RET(GetTilingDataDevMem(runInfo, &tilingDataDevPtr,\
        (crackAddr.size() + 1) * MEMSET_CRACK_TILING_DATA_MAX_SZIE));
    tilingDataLock.unlock();

    // 构造桩函数
    std::string opName;
    CHK_RET(GetOpProperty(opDescInfo, "binFileName", opName));

    s8 *stubFunc = nullptr;
    string kernelName = "";
    CHK_RET(GetOpProperty(opDescInfo, "kernelName", kernelName));

    std::unique_lock<std::mutex> stubFuncLock(stubFuncMutex_);
    CHK_RET(GetstubFunc(opName, stubFunc, kernelName));
    stubFuncLock.unlock();

    // 执行OP
    CHK_RET(
        ExecuteKernelLaunch(crackAddr, stream, reinterpret_cast<void *>(stubFunc), runInfo, tilingDataDevPtr));
    return HCCL_SUCCESS;
}
HcclResult TbeCrackCleard::ExecuteKernelLaunch(std::vector<std::int64_t> &crackAddr, void *stream,
    const void *stubFunc, OpRunInfo &runInfo, void *tilingDataDevPtr)
{
    std::unique_lock<std::mutex> crackDataLock(crackDataMutex_);

    u64 addrListDevSize = ((crackAddr.size()) * sizeof(std::int64_t) + MEMSET_ALIGNMENT_CRACK_LIST_MAX_SZIE - 1) \
        * MEMSET_ALIGNMENT_CRACK_LIST_MAX_SZIE / MEMSET_ALIGNMENT_CRACK_LIST_MAX_SZIE;

    CHK_RET(hrtMalloc(&addrListDevPtr_, addrListDevSize));
    CHK_PTR_NULL(addrListDevPtr_);
    HcclResult ret = hrtMemcpy(reinterpret_cast<void *>(static_cast<s8 *>(addrListDevPtr_)),\
        crackAddr.size() * sizeof(std::int64_t), crackAddr.data(), crackAddr.size() * sizeof(std::int64_t),\
        HcclRtMemcpyKind::HCCL_RT_MEMCPY_KIND_HOST_TO_DEVICE);
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("[ExecuteKernelLaunch]crackAddr data memcpy host to device failed"), HCCL_E_RUNTIME);
    crackDataLock.unlock();

    std::vector<const void *> deviceAddrs;
    deviceAddrs.push_back(addrListDevPtr_);
    deviceAddrs.push_back(tilingDataDevPtr);

    ret = hrtKernelLaunch(stubFunc, runInfo.blockDim, deviceAddrs.data(),
        deviceAddrs.size() * sizeof(void *), nullptr, stream);
    CHK_PRT_RET(ret != HCCL_SUCCESS,
        HCCL_ERROR("[Execute][KernelLaunch]execute kernel launch failed"), HCCL_E_RUNTIME);
    return HCCL_SUCCESS;
}

HcclResult TbeCrackCleard::GetopInfoIndex(std::string &opInfoIndex)
{
    // 芯片类型适配
    switch (deviceType_) {
        default:
        case DevType::DEV_TYPE_910:
            opInfoIndex += "910";
            break;
        case DevType::DEV_TYPE_910B:
            opInfoIndex += "910b";
            break;
        case DevType::DEV_TYPE_310P3:
            opInfoIndex += "310p3";
            break;
        case DevType::DEV_TYPE_910_93:
            break;
        case DevType::DEV_TYPE_310P1:
            HCCL_ERROR("[Get][OpBinaryPath]devType[%u] is not supported", deviceType_);
            return HCCL_E_NOT_SUPPORT;
    }

    return HCCL_SUCCESS;
}

HcclResult TbeCrackCleard::GetOpInfo(nlohmann::json &opDescInfo, nlohmann::json &opTilingInfo)
{
    std::string opInfoIndex = "MemSet_dynamic_AtomicAddrClean_1_ascend"; // TBE支持下发清零算子名称
    CHK_RET(GetopInfoIndex(opInfoIndex));
    if (opInfoMap_.find(opInfoIndex) == opInfoMap_.end()) {
        HCCL_ERROR("[Get][OpInfo]can't find tbereduce info,opInfoIndex[%s]", opInfoIndex.c_str());
        return HCCL_E_UNAVAIL;
    }
    opDescInfo = opInfoMap_[opInfoIndex];

    opInfoIndex += "_tiling_info";
    if (opInfoMap_.find(opInfoIndex) == opInfoMap_.end()) {
        HCCL_ERROR("[Get][OpInfo]can't find tbereduce info,opInfoTilingIndex[%s]", opInfoIndex.c_str());
        return HCCL_E_UNAVAIL;
    }
    opTilingInfo = opInfoMap_[opInfoIndex];

    return HCCL_SUCCESS;
}

void TbeCrackCleard::InitOpInfoMap910A()
{
    opInfoMap_["MemSet_dynamic_AtomicAddrClean_1_ascend910"] = MemSet_dynamic_AtomicAddrClean_1_ascend910;
    opInfoMap_["MemSet_dynamic_AtomicAddrClean_1_ascend910_tiling_info"] =\
        MemSet_dynamic_AtomicAddrClean_1_ascend910_tiling_info;
}

void TbeCrackCleard::InitOpInfoMap310P3()
{
    opInfoMap_["MemSet_dynamic_AtomicAddrClean_1_ascend310p3"] = MemSet_dynamic_AtomicAddrClean_1_ascend310p3;
    opInfoMap_["MemSet_dynamic_AtomicAddrClean_1_ascend310p3_tiling_info"] =\
        MemSet_dynamic_AtomicAddrClean_1_ascend310p3_tiling_info;
}

void TbeCrackCleard::InitOpInfoMap910B()
{
    opInfoMap_["MemSet_dynamic_AtomicAddrClean_1_ascend910b"] = MemSet_dynamic_AtomicAddrClean_1_ascend910b;
    opInfoMap_["MemSet_dynamic_AtomicAddrClean_1_ascend910b_tiling_info"] =\
        MemSet_dynamic_AtomicAddrClean_1_ascend910b_tiling_info;
}
} // namespace TbeReduce